package org.example.myleet.p677;

public class MapSum {

    TrieTree root;

    public MapSum() {
        root = new TrieTree();
    }

    public void insert(String key, int val) {
        root.insert(key, val);
    }

    public int sum(String prefix) {
        return root.getSum(prefix);
    }

    static class TrieNode {
        boolean end;
        int value;
        int sum;
        public TrieNode[] children;

        public TrieNode() {
            sum = 0;
            children = new TrieNode[26];
        }
    }

    static class TrieTree {
        public TrieNode[] root;

        public TrieTree() {
            root = new TrieNode[26];
        }

        public void insert(String key, int value) {
            insertSearch(root, key, 0, value);
        }

        /**
         * DFS插入字典树
         * @param curLevel 当前字典树层级
         * @param key key
         * @param i key的第i个字符
         * @param value value
         * @return 如果找到已经存在的key映射的value，则返回替换之前的value，否则返回-1
         */
        private int insertSearch(TrieNode[] curLevel, String key, int i, int value) {
            char c = key.charAt(i);
            TrieNode t = curLevel[c - 'a'];
            if (null == t) {
                t = new TrieNode();
                curLevel[c - 'a'] = t;
            }
            int prevValue;
            //叠加sum
            t.sum += value;
            if (i == key.length() - 1) {
                //到叶子节点，如果已经存在映射，就返回已经存在的值（prevValue），并替换成最新的value，并标记一个key映射的存在
                prevValue = t.end ? t.value : -1;
                t.value = value;
                t.end = true;
            } else {
                //未到叶子节点，继续深入
                prevValue = insertSearch(t.children, key, i + 1, value);
            }
            if (prevValue > 0) {
                //存在旧的映射，则sum减去旧映射，维护正确的前缀sum值
                t.sum -= prevValue;
            }
            return prevValue;
        }

        public int getSum(String prefix) {
            TrieNode[] curLevel = root;
            TrieNode t = null;
            for (char c : prefix.toCharArray()) {
                t = curLevel[c - 'a'];
                if (null == t) {
                    return 0;
                }
                curLevel = t.children;
            }
            if (null == t) {
                return 0;
            }
            return t.sum;
        }
    }
}
